from funasr import AutoModel
import soundfile
import os
import glob
import re

def check_vad(vad_list):
    # 检查列表长度是否是2的整数倍
    if len(vad_list) % 2 == 0:
        # 获取最后两个子列表
        last_two = vad_list[-2:]
        # 分别取每个子列表的最大值
        max_values = [max(sublist) for sublist in last_two]
        # 计算两个最大值之间的绝对差值
        result = abs(max_values[1] - max_values[0])
        return result
    else:
        # 如果不是2的整数倍，可以抛出异常或返回特定值表示错误
        return 0

def natural_keys(text):
    """
    将文件名解析为包含整数的部分，以实现自然排序。
    """
    return [int(c) if c.isdigit() else c for c in re.split(r'(\d+)', text)]


# 表示单个音频的长度.
chunk_size = 60  # ms
model = AutoModel(model="/root/yuehu/assets/speech_fsmn_vad_zh-cn-16k-common-pytorch", device='cuda:1',disable_update=True)

# 目标文件夹路径
directory_path = 'app/input/f83ef92d8f4544969a35bd40d0bf0bb1'  # 将此替换为实际的目录路径

# 使用glob获取所有.wav文件
wav_files = glob.glob(os.path.join(directory_path, '*'))
wav_files.sort(key=natural_keys)

cache = {}
vad_res = []

for wav_file in wav_files:
    print(f"Processing file: {wav_file}")
    speech, sample_rate = soundfile.read(wav_file)
    chunk_stride = int(chunk_size * sample_rate / 1000)
    total_chunk_num = int((len(speech) - 1) / chunk_stride + 1)

    # print(f'sample_rate: {sample_rate}, chunk_stride: {chunk_stride}, speech len: {len(speech)}, total_chunk_num: {total_chunk_num}')
    for i in range(total_chunk_num):
        speech_chunk = speech[i*chunk_stride:(i+1)*chunk_stride]
        # is_final = i == total_chunk_num - 1
        is_final = False
        res = model.generate(input=speech_chunk, disable_update=True,output_dir='output_dir/', is_final=is_final, chunk_size=chunk_size, max_end_silence_time=400, max_start_silence_time=1000)
        if len(res) == 0:
            continue
        value = res[0]['value']
        if len(value) == 0:
            continue
        vad_res.append(res[0]['value'][0])

        # 表示至少识别 1 秒的长度才认为需要结束.
        if check_vad(vad_res) > 1000:
            print(f'-----------------------------------')
            break
    else:
        # 如果内层循环没有遇到break，则继续处理下一个文件
        continue

print(vad_res)
